import pymysql
from dbutils.pooled_db import PooledDB
from datetime import datetime
import os
from db_res import dbs

"""
数据库备份，目前仅支持mysql。
"""


class DbBackup:
    def __init__(self, tag):
        self.tag = tag
        self.init_db_pool()

    def init_db_pool(self):
        self.pool = {}
        for db in dbs:
            self.pool[db] = PooledDB(
                # 使用链接数据库的模块
                creator=pymysql,
                # 连接池允许的最大连接数，0和None表示不限制连接数
                maxconnections=6,
                # 初始化时，链接池中至少创建的空闲的链接，0表示不创建
                mincached=2,
                # 链接池中最多闲置的链接，0和None不限制
                maxcached=5,
                # 链接池中最多共享的链接数量，0和None表示全部共享。PS: 无用，因为pymysql和MySQLdb等模块的 threadsafety都为1，所有值无论设置为多少，_maxcached永远为0，所以永远是所有链接都共享。
                maxshared=3,
                # 连接池中如果没有可用连接后，是否阻塞等待。True，等待；False，不等待然后报错
                blocking=True,
                # 一个链接最多被重复使用的次数，None表示无限制
                maxusage=None,
                # 开始会话前执行的命令列表。如：["set datestyle to ...", "set time zone ..."]
                setsession=[],
                # ping MySQL服务端，检查是否服务可用。# 如：0 = None = never, 1 = default = whenever it is requested, 2 = when a cursor is created, 4 = when a query is executed, 7 = always
                ping=0,
                host=dbs[db]["ip"],
                port=dbs[db]["port"],
                user=dbs[db]["usr"],
                password=dbs[db]["pwd"],
                database=dbs[db]["db"],
                charset=dbs[db]["charset"]
            )

    def exec(self, dbName, sql, cs=""):
        conn = self.pool[dbName].connection()
        cs = pymysql.cursors.DictCursor if cs != None else cs
        cursor = conn.cursor(cursor=cs)
        try:
            cursor.execute(sql)
            data = cursor.fetchall()
        except Exception as e:
            print(e)
        finally:
            cursor.close()
            conn.close()
        return data

    def show_tables(self, dbName):
        result = []
        for i in self.exec(dbName, "show tables;", cs=None):
            for j in i:
                result.append(j)
        return result

    def run(self):
        for i in self.pool:
            path = f"{i}_{self.tag}"
            if not os.path.exists(path):
                os.makedirs(path)
            for j in self.show_tables(i):
                self.deal_table(i, j)

    def deal_table(self, dbName, tblName):
        cnt = self.exec(dbName, f"select count(1) as cnt from {tblName};")[0]["cnt"]
        batch_size = 1000
        idx = 0
        while idx < cnt:
            data = self.exec(dbName, f"select * from {tblName} limit {idx}, {batch_size};")
            data = dict_insert_sql(tblName, data)
            write_file(f"{dbName}_{self.tag}/{tblName}_{self.tag}.sql", data)
            now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            log_progress(f"[{now}] {dbName} {tblName} {idx}")
            idx += batch_size


def dict_insert_sql(table, data):
    len_d = len(data)
    if len_d <= 0: return ""
    fields = [i for i in data[0]]
    header = ", ".join(fields)
    sql = f"insert into {table} ({header}) values \n"
    values = []
    for i in data:
        value = ""
        for j in fields:
            v = i[j]
            if v is None:
                value += ", null"
            elif type(v) is int or type(v) is float:
                value += f", {v}"
            else:
                if type(v) is str:
                    v = v.replace("'", "\\'")
                value += f", '{v}'"
        values.append(f"({value[2:]})")
    sql += ",\n".join(values)
    return sql + ";"


def write_file(filename, data):
    with open(filename, "a+") as f:
        f.write(data + "\n")


def log_progress(data):
    print("\r" + " " * 100, end="")
    print(f"\r{data}", end="")


if __name__ == "__main__":
    begin = datetime.now()
    begin_str = begin.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{begin_str}] 开始备份...")
    DbBackup("2024-01-28").run()
    end = datetime.now()
    end_str = end.strftime("%Y-%m-%d %H:%M:%S")
    delt = end - begin
    print(f"\n[{end_str}] 用时: " + str(delt.seconds) + "." + str(delt.microseconds // 100000) + " 秒")
